/*
 * Copyright 2011-2025 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0.
 * See `LICENSE` in the project root for license information.
 */

package me.ijleex.platform.system.authn.util;

import java.time.Duration;
import java.util.AbstractMap;
import java.util.Map;
import java.util.function.Function;

import org.apache.commons.lang3.StringUtils;
import org.jose4j.jwk.EcJwkGenerator;
import org.jose4j.jwk.JsonWebKey;
import org.jose4j.jwk.PublicJsonWebKey;
import org.jose4j.jwk.Use;
import org.jose4j.jws.AlgorithmIdentifiers;
import org.jose4j.jws.JsonWebSignature;
import org.jose4j.jwt.JwtClaims;
import org.jose4j.jwt.MalformedClaimException;
import org.jose4j.jwt.NumericDate;
import org.jose4j.jwt.consumer.InvalidJwtException;
import org.jose4j.jwt.consumer.JwtConsumer;
import org.jose4j.jwt.consumer.JwtConsumerBuilder;
import org.jose4j.keys.EllipticCurves;
import org.jose4j.lang.JoseException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import me.ijleex.platform.framework.util.snowflake.IdWorker;

/**
 * 令牌（JWT）工具.
 *
 * @author liym
 * @since 2019-12-20 21:51 新建
 */
public final class JwtUtil {

    /**
     * 生成密钥的算法
     *
     * <p>在开发应用的时候启用 JWT，使用 ES256、RS256 更加安全（推荐 ES256 算法）</p>
     *
     * @see AlgorithmIdentifiers#RSA_USING_SHA256
     */
    private static final String KEY_ALGORITHM = AlgorithmIdentifiers.ECDSA_USING_P256_CURVE_AND_SHA256;
    /**
     * Token过期时长
     */
    private static final Duration DEFAULT_EXPIRE_TIME = Duration.ofMinutes(40L);
    /**
     * <a href="https://openid.net/specs/openid-connect-core-1_0.html#SelfIssued">Self-Issued issuer</a>
     */
    private static final String SELF_ISSUED_ISSUER = "https://self-issued.me";

    private static final Logger logger = LoggerFactory.getLogger(JwtUtil.class);

    private JwtUtil() {
        throw new IllegalAccessError("Utility class");
    }

    /**
     * 生成密钥对，包含私钥和公钥
     *
     * @return 密钥对 {@code 私钥:公钥}
     * @since 2023-06-06 22:11
     */
    public static Map.Entry<String, String> createKeyPair() {
        JsonWebKey jwk;
        try {
            jwk = EcJwkGenerator.generateJwk(EllipticCurves.P256);
        } catch (JoseException e) {
            logger.error("createKeyPair error", e);
            throw new IllegalStateException("createKeyPair error", e);
        }

        final String keyId = Long.toString(IdWorker.getId());
        jwk.setUse(Use.SIGNATURE);
        jwk.setKeyId(keyId);
        jwk.setAlgorithm(KEY_ALGORITHM);

        String keyPairJson = jwk.toJson(JsonWebKey.OutputControlLevel.INCLUDE_PRIVATE);
        String publicKeyJson = jwk.toJson();

        return new AbstractMap.SimpleEntry<>(keyPairJson, publicKeyJson);
    }

    /**
     * 生成 Token，默认过期时间40分钟
     *
     * @param keyPairJson 密钥对
     * @param loginId 登录名，如{@code ijliym}
     * @param clientId 客户端ID
     * @return Token
     * @see #DEFAULT_EXPIRE_TIME
     * @since 2023-06-06 23:08
     */
    public static String createToken(String keyPairJson, String loginId, String clientId) {
        return createToken(keyPairJson, loginId, clientId, DEFAULT_EXPIRE_TIME);
    }

    /**
     * 生成 Token
     *
     * @param keyPairJson 密钥对
     * @param loginId 登录名，如{@code ijliym}
     * @param clientId 客户端ID
     * @param expireTime Token过期时长
     * @return Token
     * @since 2019-12-20 21:54 创建方法
     * @since 2023-06-06 22:33 实现方法，添加参数
     */
    public static String createToken(String keyPairJson, String loginId, String clientId, Duration expireTime) {
        String idToken;
        try {
            PublicJsonWebKey keyPairJwk = PublicJsonWebKey.Factory.newPublicJwk(keyPairJson);
            JsonWebSignature jws = getJsonWebSignature(keyPairJwk, SELF_ISSUED_ISSUER, loginId, clientId, expireTime);

            // signing token
            idToken = jws.getCompactSerialization();
        } catch (JoseException e) {
            logger.error("createToken error", e);
            throw new IllegalStateException("createToken error", e);
        }
        return idToken;
    }

    /**
     * 初始化{@link JsonWebSignature}，用于生成Token
     *
     * @param keyPairJwk 私钥
     * @param issuer JWT签发者
     * @param subject JWT所面向的用户
     * @param audience Token的受众
     * @param expireTime Token过期时长
     * @return {@link JsonWebSignature}
     * @since 2023-06-06 22:45
     */
    private static JsonWebSignature getJsonWebSignature(PublicJsonWebKey keyPairJwk,
            String issuer, String subject, String audience, Duration expireTime) {
        final String tokenId = Long.toString(IdWorker.getId());

        // Payload
        JwtClaims claims = new JwtClaims();
        if (StringUtils.isNotBlank(issuer)) {
            claims.setIssuer(issuer); // https://...
        }
        if (StringUtils.isNotBlank(subject)) {
            claims.setSubject(subject); // loginId
        }
        if (StringUtils.isNotBlank(audience)) {
            claims.setAudience(audience); // client_id
        }
        claims.setExpirationTimeMinutesInTheFuture(expireTime.toMinutes());
        claims.setIssuedAtToNow();
        claims.setNotBeforeMinutesInThePast(2.0F);
        claims.setJwtId(tokenId);

        // JWS
        JsonWebSignature jws = new JsonWebSignature();
        jws.setHeader("typ", "JWT");
        jws.setPayload(claims.toJson());

        // 签名算法
        jws.setAlgorithmHeaderValue(KEY_ALGORITHM);
        jws.setKeyIdHeaderValue(keyPairJwk.getKeyId());
        jws.setKey(keyPairJwk.getPrivateKey());

        return jws;
    }

    /**
     * 解码验证 Token，判断是否合法、是否过期、…
     *
     * @param publicKeyJson 公钥
     * @param loginId 登录名，如{@code ijliym}
     * @param clientId 客户端ID
     * @param expireTime Token过期时长
     * @param idToken Token
     * @return true/false
     * @since 2019-12-20 23:48 创建方法
     * @since 2023-06-06 23:11 实现方法，添加参数
     */
    public static boolean isTokenValid(String publicKeyJson,
            String loginId, String clientId, Duration expireTime,
            String idToken) {
        boolean valid = false;
        try {
            PublicJsonWebKey publicKey = PublicJsonWebKey.Factory.newPublicJwk(publicKeyJson);
            JwtConsumer consumer = getJwtConsumer(publicKey, SELF_ISSUED_ISSUER, loginId, clientId, expireTime);

            // {"sub":"admin","aud":"yjfxlxt_app","nbf":1684934965,"exp":1684938685,"iat":1684935085,"jti":"557819634012680192"}
            JwtClaims claims = consumer.processToClaims(idToken);
            valid = claims != null;
        } catch (JoseException e) {
            logger.error("isTokenValid error；公钥不合法={}", publicKeyJson, e);
        } catch (InvalidJwtException e) {
            logger.error("isTokenValid error；Token校验不通过={}", idToken, e);
        }
        return valid;
    }

    /**
     * 初始化{@link JwtConsumer}，用于解码Token
     *
     * @param publicKey 公钥
     * @param issuer JWT签发者
     * @param subject JWT所面向的用户
     * @param audience Token的受众
     * @param expireTime Token过期时长
     * @return {@link JwtConsumer}
     * @since 2023-06-06 23:28
     */
    private static JwtConsumer getJwtConsumer(PublicJsonWebKey publicKey,
            String issuer, String subject, String audience, Duration expireTime) {
        JwtConsumerBuilder builder = new JwtConsumerBuilder().setVerificationKey(publicKey.getPublicKey());
        if (StringUtils.isNotBlank(issuer)) {
            builder.setExpectedIssuer(true, issuer); // https://...
        }
        if (StringUtils.isNotBlank(subject)) {
            builder.setExpectedSubject(subject); // loginId
        }
        if (StringUtils.isNotBlank(audience)) {
            builder.setExpectedAudience(true, audience); // client_id
        }

        builder.setRequireJwtId()
                .setRequireExpirationTime()
                .setRequireIssuedAt()
                .setRequireNotBefore()
                .setAllowedClockSkewInSeconds(10 * 60) // 允许系统时钟偏离时间（10分钟）
                .setMaxFutureValidityInMinutes((int) expireTime.toMinutes()); // Token有效期

        return builder.build();
    }

    /**
     * 解码验证 Token，并获取登录名
     *
     * @param publicKeyJson 公钥
     * @param clientId 客户端ID
     * @param expireTime Token过期时长
     * @param idToken Token
     * @return 登录名（loginId）
     * @since 2023-06-07 21:44
     */
    public static String getLoginIdFromToken(final String publicKeyJson,
            String clientId, Duration expireTime, String idToken) {
        return extractClaim(publicKeyJson, SELF_ISSUED_ISSUER, null, clientId, expireTime, idToken,
                jwtClaims -> {
                    try {
                        return jwtClaims.getSubject();
                    } catch (MalformedClaimException e) {
                        logger.error("getLoginIdFromToken error", e);
                        return null;
                    }
                });
    }

    /**
     * 解码验证 Token，并获取过期时间
     *
     * @param publicKeyJson 公钥
     * @param clientId 客户端ID
     * @param expireTime Token过期时长
     * @param idToken Token
     * @return Token过期时间（时间戳，毫秒）
     * @since 2025-06-26 20:42
     */
    public static Long getExpireTimeFromToken(final String publicKeyJson,
            String clientId, Duration expireTime, String idToken) {
        return extractClaim(publicKeyJson, SELF_ISSUED_ISSUER, null, clientId, expireTime, idToken,
                jwtClaims -> {
                    try {
                        NumericDate expireTimeSec = jwtClaims.getExpirationTime();
                        return expireTimeSec.getValueInMillis();
                    } catch (MalformedClaimException e) {
                        logger.error("getExpireTimeFromToken error", e);
                        return null;
                    }
                });
    }

    /**
     * @since 2023-06-07 00:03
     */
    private static <T> T extractClaim(final String publicKeyJson,
            String issuer, String subject, String audience, Duration expireTime,
            String idToken, Function<JwtClaims, T> claimsResolver) {
        final JwtClaims claims = extractAllClaims(publicKeyJson, issuer, subject, audience, expireTime, idToken);
        if (claims == null) {
            return null;
        }
        return claimsResolver.apply(claims);
    }

    /**
     * 解码验证Token，获取生成Token的属性
     */
    private static JwtClaims extractAllClaims(String publicKeyJson,
            String issuer, String subject, String audience, Duration expireTime,
            String idToken) {
        JwtClaims claims = null;
        try {
            PublicJsonWebKey publicKey = PublicJsonWebKey.Factory.newPublicJwk(publicKeyJson);
            JwtConsumer consumer = getJwtConsumer(publicKey, issuer, subject, audience, expireTime);

            // {"sub":"admin","aud":"yjfxlxt_app","nbf":1684934965,"exp":1684938685,"iat":1684935085,"jti":"557819634012680192"}
            claims = consumer.processToClaims(idToken);
        } catch (JoseException e) {
            logger.error("extractAllClaims error；公钥不合法={}", publicKeyJson, e);
        } catch (InvalidJwtException e) {
            logger.error("extractAllClaims error；Token校验不通过={}", idToken, e);
        }
        return claims;
    }

}
